#!/usr/bin/env python
# coding: utf-8

"""
dataset_splite.py
create by cuiyu 2025-02-20
说明：
 * 使用 ultralytics 工具自动划分数据集
 * 不会更改输入目录下的任何内容，
 * 在输入目录的 ../ 中创建 autosplit_train/val/test.txt 文件，分别保存对应的图片全路径名

用法:
  dataset_splite.py

更新:
2025-02-20:
"""
from pathlib import Path
from ultralytics.data.utils import autosplit
# from ultralytics.data.split import autosplit


def datasetSplit(data_root: Path):
    """
    自动划分数据集到训练集、验证集和测试集
    
    Args:
        :data_root: str, 数据集根目录路径
    """
    autosplit(  
        path=data_root,
        weights=(0.9, 0.1, 0.0),  # (train, validation, test) fractional splits
        annotated_only=True,  # True: 只有对应的 .txt 标注文件存在的图片才会使用; False：使用所有图片
    )
    # 自动生成的文件保存在 data_root 的父目录中
    # 自动生成文件：autosplit_train.txt / autosplit_val.txt / autosplit_test.txt
    
    # 修改文件名，删除 autosplit_ 前缀
    parent_dir = Path(data_root).parent
    for prefix in ["train", "val", "test"]:
        src = parent_dir / f"autosplit_{prefix}.txt"
        dest = parent_dir / f"{prefix}.txt"
        if src.exists():
            if dest.exists():
                dest.unlink()
            src.rename(dest)
    pass
pass

def _main():
    # data_root = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/rddhw/label/20250111")
    data_root = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/井盖病害/训练数据/img01")
    datasetSplit(data_root)
    pass

if __name__ == "__main__":
    _main()